package pkg

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"fmt"
	"math/rand"
	"strings"
	"time"
)

const AES_HEADER = "AES_"

var aesLetter = []rune("0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

type AESEncryptor struct {
}

type aesEncodeValue struct {
	cipherText []byte
}

// NewAESEncodeValue创建新的AES加密后得值
func NewAESEncodeValue(cipherText []byte) *aesEncodeValue {
	return &aesEncodeValue{
		cipherText: cipherText,
	}
}

// EncodeValue 返回一个有特征得AES得加密值
func (e *aesEncodeValue) EncodeValue() string {
	encodeText := base64.StdEncoding.EncodeToString(e.cipherText)
	return fmt.Sprintf("%s%s", AES_HEADER, encodeText)
}

type aesDecodeValue struct {
	cipherText []byte
}

// NewAESDecodeValue创建新的AES解密后得值
func NewAESDecodeValue(cipherText []byte) *aesDecodeValue {
	return &aesDecodeValue{cipherText: cipherText}
}

// DecodeValue 返回一个去特征得AES解密值
func (d *aesDecodeValue) DecodeValue() ([]byte, error) {
	if !strings.HasPrefix(string(d.cipherText), AES_HEADER) {
		return d.cipherText, nil
	}
	return base64.StdEncoding.DecodeString(
		strings.TrimPrefix(string(d.cipherText), AES_HEADER))
}

// IsAESValue 判断是否为aes加密后的值
func (d *aesDecodeValue) IsAESValue() bool {
	return strings.HasPrefix(string(d.cipherText), AES_HEADER)
}
func pkcs7Padding(cipherText []byte, blockSize int) []byte {
	padding := blockSize - len(cipherText)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(cipherText, padText...)
}
func pkcs7UnPadding(plantText []byte, blockSize int) []byte {
	length := len(plantText)
	unPadding := int(plantText[length-1])
	return plantText[:(length - unPadding)]
}

// Encrypt 用于aesKey 加密
func (e AESEncryptor) Encrypt(aesKey, content string) (string, error) {
	block, err := aes.NewCipher([]byte(aesKey))
	if err != nil {
		return "", err
	}
	plantText := pkcs7Padding([]byte(content), block.BlockSize())
	blockModel := cipher.NewCBCEncrypter(block, []byte(aesKey)[:block.BlockSize()])
	cipherText := make([]byte, len(plantText))
	blockModel.CryptBlocks(cipherText, plantText)
	encode := NewAESEncodeValue(cipherText)
	return encode.EncodeValue(), nil
}

// Decrypt 用于aesKey 解密
func (e AESEncryptor) Decrypt(aesKey string, content string) (string, error) {
	decode := NewAESDecodeValue([]byte(content))
	//如果不是aes 加密后的内容 直接返回
	if !decode.IsAESValue() {
		return content, nil
	}
	contentBytes, err := decode.DecodeValue()
	if err != nil {
		return "", err
	}
	block, err := aes.NewCipher([]byte(aesKey))
	if err != nil {
		return "", err
	}
	blockModel := cipher.NewCBCDecrypter(block, []byte(aesKey)[:block.BlockSize()])
	plantText := make([]byte, len(contentBytes))
	blockModel.CryptBlocks(plantText, contentBytes)
	plantText = pkcs7UnPadding(plantText, block.BlockSize())
	return string(plantText), nil
}

// RandKey 随机生成32位  aes key
func (e AESEncryptor) RandKey() string {
	b := make([]rune, 32)
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	for i := range b {
		b[i] = aesLetter[r.Intn(62)]
	}
	return string(b)
}
